#pip install numba

import numpy as np
import matplotlib.pyplot as plt
from numba import njit

@njit
def mandelbrot(c, max_iter):
    z = 0
    for n in range(max_iter):
        if abs(z) > 2:
            return n
        z = z*z + c
    return max_iter

@njit
def mandelbrot_set(xmin,xmax,ymin,ymax,width,height,max_iter):
    print("Génération de l'image en cours...")
    r1 = np.linspace(xmin, xmax, width)
    r2 = np.linspace(ymin, ymax, height)
    n3 = np.empty((width,height))
    for i in range(width):
        for j in range(height):
            n3[j,i] = mandelbrot(r1[i] + 1j * r2[j], max_iter)  # Note the index order
    print("Image générée.")
    print()
    return (r1,r2,n3)

def on_click(event):
    if event.inaxes == ax:
        x, y = event.xdata, event.ydata
        width = abs(ax.get_xlim()[1] - ax.get_xlim()[0])
        height = abs(ax.get_ylim()[1] - ax.get_ylim()[0])
        ax.set_xlim(x - width/10, x + width/10)
        ax.set_ylim(y - height/10, y + height/10)
        d = mandelbrot_set(x - width/10, x + width/10, y - height/10, y + height/10, 1000, 1000, 256)
        image.set_data(d[2])
        image.set_extent((x - width/10, x + width/10, y - height/10, y + height/10))
        plt.draw()

fig, ax = plt.subplots()

xmin, xmax, ymin, ymax = -2.0, 1.0, -1.5, 1.5
d = mandelbrot_set(xmin, xmax, ymin, ymax, 1000, 1000, 256)
image = ax.imshow(d[2], extent=(xmin, xmax, ymin, ymax), cmap='turbo', aspect='auto', origin='lower')

fig.canvas.mpl_connect('button_press_event', on_click)

plt.show()